from functools import wraps
from flask import Flask, jsonify, Response, request, redirect, url_for, send_from_directory
import flask
import os
import uuid
import pandas as pd

# Custom QwenLLM implementation
from vanna.base import VannaBase
from vanna.chromadb import ChromaDB_VectorStore
from dashscope import Generation
import random

# Memory Cache for storing session data
class MemoryCache:
    def __init__(self):
        self.cache = {}

    def generate_id(self, *args, **kwargs):
        return str(uuid.uuid4())

    def set(self, id, field, value):
        if id not in self.cache:
            self.cache[id] = {}

        self.cache[id][field] = value

    def get(self, id, field):
        if id not in self.cache:
            return None

        if field not in self.cache[id]:
            return None

        return self.cache[id][field]

    def get_all(self, field_list) -> list:
        return [
            {
                "id": id,
                **{
                    field: self.get(id=id, field=field)
                    for field in field_list
                }
            }
            for id in self.cache
        ]

    def delete(self, id):
        if id in self.cache:
            del self.cache[id]

# QwenLLM Class
class QwenLLM(VannaBase):
    def __init__(self, config=None):
        self.model = config['model']
        self.api_key = config['api_key']
  
    def system_message(self, message: str):
        return {'role': 'system', 'content': message}

    def user_message(self, message: str):
        return {'role': 'user', 'content': message}

    def assistant_message(self, message: str):
        return {'role': 'assistant', 'content': message}
  
    def submit_prompt(self, prompt, **kwargs):
        resp = Generation.call(
            model=self.model,
            messages=prompt,
            seed=random.randint(1, 10000),
            result_format='message',
            api_key=self.api_key)
        answer = resp.output.choices[0].message.content
        return answer

# Combined VannaBase class
class MyVanna(ChromaDB_VectorStore, QwenLLM):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        QwenLLM.__init__(self, config=config)

# Initialize Flask app
app = Flask(__name__, static_url_path='')

# Initialize cache
cache = MemoryCache()

# Initialize Vanna with QwenLLM
vn = MyVanna({'api_key': 'sk-d8d15e769aa549129467e605c1f2905e', 'model': 'qwen-max'})

# Connect to MySQL DB
try:
    vn.connect_to_mysql(host='172.16.98.128', dbname='volbeacon', user='chinaedu', password='chinaedu', port=13002)
except Exception as e:
    print(f"Warning: Could not connect to MySQL: {str(e)}")
    print("You will need to connect to a database before using the application.")

# Helper decorator for cache requirements
def requires_cache(fields):
    def decorator(f):
        @wraps(f)
        def decorated(*args, **kwargs):
            id = request.args.get('id')

            if id is None:
                return jsonify({"type": "error", "error": "No id provided"})
            
            for field in fields:
                if cache.get(id=id, field=field) is None:
                    return jsonify({"type": "error", "error": f"No {field} found"})
            
            field_values = {field: cache.get(id=id, field=field) for field in fields}
            
            # Add the id to the field_values
            field_values['id'] = id

            return f(*args, **field_values, **kwargs)
        return decorated
    return decorator

# API Routes
@app.route('/api/v0/generate_questions', methods=['GET'])
def generate_questions():
    try:
        questions = vn.generate_questions()
        return jsonify({
            "type": "question_list", 
            "questions": questions,
            "header": "Here are some questions you can ask:"
        })
    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/generate_sql', methods=['GET'])
def generate_sql():
    question = flask.request.args.get('question')

    if question is None:
        return jsonify({"type": "error", "error": "No question provided"})

    id = cache.generate_id(question=question)
    
    try:
        sql = vn.generate_sql(question=question)

        cache.set(id=id, field='question', value=question)
        cache.set(id=id, field='sql', value=sql)

        return jsonify(
            {
                "type": "sql", 
                "id": id,
                "text": sql,
            })
    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/run_sql', methods=['GET'])
@requires_cache(['sql'])
def run_sql(id: str, sql: str):
    try:
        df = vn.run_sql(sql=sql)

        cache.set(id=id, field='df', value=df)

        return jsonify(
            {
                "type": "df", 
                "id": id,
                "df": df.head(10).to_json(orient='records'),
            })

    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/download_csv', methods=['GET'])
@requires_cache(['df'])
def download_csv(id: str, df):
    csv = df.to_csv()

    return Response(
        csv,
        mimetype="text/csv",
        headers={"Content-disposition":
                 f"attachment; filename={id}.csv"})

@app.route('/api/v0/generate_plotly_figure', methods=['GET'])
@requires_cache(['df', 'question', 'sql'])
def generate_plotly_figure(id: str, df, question, sql):
    try:
        code = vn.generate_plotly_code(question=question, sql=sql, df_metadata=f"Running df.dtypes gives:\n {df.dtypes}")
        fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
        fig_json = fig.to_json()

        cache.set(id=id, field='fig_json', value=fig_json)

        return jsonify(
            {
                "type": "plotly_figure", 
                "id": id,
                "fig": fig_json,
            })
    except Exception as e:
        # Print the stack trace
        import traceback
        traceback.print_exc()

        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/get_training_data', methods=['GET'])
def get_training_data():
    try:
        df = vn.get_training_data()
        
        # If df is None or empty, return empty dataframe
        if df is None or len(df) == 0:
            return jsonify({
                "type": "df", 
                "id": "training_data",
                "df": pd.DataFrame().to_json(orient='records'),
            })
        
        return jsonify({
            "type": "df", 
            "id": "training_data",
            "df": df.head(25).to_json(orient='records'),
        })
    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/remove_training_data', methods=['POST'])
def remove_training_data():
    # Get id from the JSON body
    id = flask.request.json.get('id')

    if id is None:
        return jsonify({"type": "error", "error": "No id provided"})

    try:
        if vn.remove_training_data(id=id):
            return jsonify({"success": True})
        else:
            return jsonify({"type": "error", "error": "Couldn't remove training data"})
    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/train', methods=['POST'])
def add_training_data():
    question = flask.request.json.get('question')
    sql = flask.request.json.get('sql')
    ddl = flask.request.json.get('ddl')
    documentation = flask.request.json.get('documentation')

    try:
        id = vn.train(question=question, sql=sql, ddl=ddl, documentation=documentation)
        return jsonify({"id": id})
    except Exception as e:
        print("TRAINING ERROR", e)
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/generate_followup_questions', methods=['GET'])
@requires_cache(['df', 'question', 'sql'])
def generate_followup_questions(id: str, df, question, sql):
    try:
        followup_questions = vn.generate_followup_questions(question=question, sql=sql, df=df)

        cache.set(id=id, field='followup_questions', value=followup_questions)

        return jsonify(
            {
                "type": "question_list", 
                "id": id,
                "questions": followup_questions,
                "header": "Here are some followup questions you can ask:"
            })
    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/load_question', methods=['GET'])
@requires_cache(['question', 'sql'])
def load_question(id: str, question, sql):
    try:
        df = cache.get(id=id, field='df')
        fig_json = cache.get(id=id, field='fig_json')
        followup_questions = cache.get(id=id, field='followup_questions')
        
        result = {
            "type": "question_cache", 
            "id": id,
            "question": question,
            "sql": sql,
        }
        
        if df is not None:
            result["df"] = df.head(10).to_json(orient='records')
        
        if fig_json is not None:
            result["fig"] = fig_json
            
        if followup_questions is not None:
            result["followup_questions"] = followup_questions
            
        return jsonify(result)

    except Exception as e:
        return jsonify({"type": "error", "error": str(e)})

@app.route('/api/v0/get_question_history', methods=['GET'])
def get_question_history():
    limit = request.args.get('limit')
    all_items = request.args.get('all', 'false').lower() == 'true'
    
    # Get all questions from cache
    questions = cache.get_all(field_list=['question'])
    
    # Only limit if specified and not requesting all
    if limit is not None and not all_items:
        try:
            limit = int(limit)
            questions = questions[:limit]
        except ValueError:
            pass  # Ignore invalid limit values
    
    return jsonify({"type": "question_history", "questions": questions})

@app.route('/', defaults={'path': ''})
@app.route('/<path:path>')
def serve_static(path):
    if path == "":
        return app.send_static_file('index.html')
    return app.send_static_file(path)

if __name__ == '__main__':
    print("Vanna App is running at http://127.0.0.1:5000")
    app.run(debug=True, host='0.0.0.0', port=5000) 